package com.deductive.netty.udp.handler;

import com.deductive.netty.udp.UdpClientManager;
import com.deductive.netty.udp.message.UdpMessage;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.socket.DatagramPacket;
import io.netty.util.CharsetUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

@Component
@ChannelHandler.Sharable
public class UdpMessageHandler extends SimpleChannelInboundHandler<DatagramPacket> {

    private static final Logger logger = LoggerFactory.getLogger(UdpMessageHandler.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    private UdpClientManager udpClientManager;

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket packet) throws Exception {
        String jsonString = packet.content().toString(CharsetUtil.UTF_8);
        InetSocketAddress sender = packet.sender();

        UdpMessage udpMessage;
        try {
            udpMessage = objectMapper.readValue(jsonString, UdpMessage.class);
        } catch (JsonProcessingException e) {
            logger.error("Failed to parse UDP JSON message from {}: {}", sender, jsonString, e);
            sendMessage(ctx, sender, createErrorResponse("Invalid JSON format"));
            return;
        }

        switch (udpMessage.getType()) {
            case "AUTH":
                handleAuth(ctx, sender, udpMessage);
                break;
            case "HEARTBEAT":
                handleHeartbeat(ctx, sender);
                break;
            case "DATA":
                handleData(ctx, sender, udpMessage);
                break;
            default:
                logger.warn("Received unknown UDP message type from {}: {}", sender, udpMessage.getType());
                sendMessage(ctx, sender, createErrorResponse("Unknown message type: " + udpMessage.getType()));
                break;
        }
    }

    private void handleAuth(ChannelHandlerContext ctx, InetSocketAddress sender, UdpMessage message) {
        JsonNode payload = message.getPayload();
        if (payload != null && payload.has("userId")) {
            String userId = payload.get("userId").asText();

            if (userId.isEmpty()) {
                sendMessage(ctx, sender, createAuthResponse(false, "User ID cannot be empty."));
                return;
            }

            udpClientManager.registerClient(userId, sender);
            logger.info("UDP user {}@{} authenticated.", userId, sender);
            sendMessage(ctx, sender, createAuthResponse(true, "Authentication successful."));
        } else {
            sendMessage(ctx, sender, createAuthResponse(false, "Authentication failed. Missing userId in payload."));
        }
    }

    private void handleHeartbeat(ChannelHandlerContext ctx, InetSocketAddress sender) {
        logger.debug("Received UDP heartbeat from {}", sender);
        // Optionally, reply to the heartbeat
        Map<String, Object> response = new HashMap<>();
        response.put("type", "HEARTBEAT_REPLY");
        response.put("payload", Collections.singletonMap("timestamp", System.currentTimeMillis()));
        sendMessage(ctx, sender, response);
    }

    private void handleData(ChannelHandlerContext ctx, InetSocketAddress sender, UdpMessage message) {
        logger.info("Received DATA message from {}: {}", sender, message.getPayload());
        // Echo the message back for now
        Map<String, Object> response = new HashMap<>();
        response.put("type", "ECHO");
        response.put("payload", message.getPayload());
        sendMessage(ctx, sender, response);
    }

    private void sendMessage(ChannelHandlerContext ctx, InetSocketAddress recipient, Object messageObject) {
        try {
            String jsonResponse = objectMapper.writeValueAsString(messageObject);
            byte[] bytes = jsonResponse.getBytes(CharsetUtil.UTF_8);
            ByteBuf buffer = Unpooled.copiedBuffer(bytes);
            ctx.writeAndFlush(new DatagramPacket(buffer, recipient));
        } catch (JsonProcessingException e) {
            logger.error("Failed to serialize response message for recipient {}", recipient, e);
        }
    }

    private Map<String, Object> createErrorResponse(String errorMessage) {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "ERROR");
        response.put("payload", Collections.singletonMap("message", errorMessage));
        return response;
    }

    private Map<String, Object> createAuthResponse(boolean success, String message) {
        Map<String, Object> payload = new HashMap<>();
        payload.put("success", success);
        payload.put("message", message);

        Map<String, Object> response = new HashMap<>();
        response.put("type", success ? "AUTH_SUCCESS" : "AUTH_FAILED");
        response.put("payload", payload);
        return response;
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        logger.error("Exception caught in UDP handler", cause);
        // For UDP, we don't close the channel as it's connectionless.
    }
}
